#include "task.h"
#include "kernel.h"
#include "status.h"
#include "memory/heap/kheap.h"
#include "memory/memory.h"
#include "process.h"
#include "idt/idt.h"
#include "memory/paging/paging.h"
#include "utils/string.h"

// The task that is currently running
struct task* current_task = 0; 

// Task linked list
struct task* task_tail = 0; 
struct task* task_head = 0; 

struct task* task_get_next()
{
    if (!current_task->next) {
        return task_head; 
    }

    return current_task->next; 
}

static void task_list_remove(struct task* task)
{
    if (task->prev) {
        task->prev->next = task->next; 
    }

    if (task == task_head) {
        task_head = task->next; 
    }

    if (task == task_tail) {
        task_tail = task->prev; 
    }

    if (task == current_task) {
        current_task = task_get_next(); 
    }
}

struct task* get_current_task()
{
    return current_task; 
}

const int task_switch(struct task* task)
{
    current_task = task; 
    paging_switch(task->page_directory); 
    return 0; 
}

const int task_page()
{
    user_registers(); 
    task_switch(current_task); 

    return 0; 
}

const int task_page_task(struct task* task) 
{
    user_registers();
    paging_switch(task->page_directory);
    return 0; 
}

void task_run_first_ever_task()
{
    if (!current_task) {
        panic("task_run_first_ever_task(): No current task exists"); 
    }

    task_switch(task_head); 
    task_return(&task_head->registers); 
}

const int task_init(struct task* task, struct process* process)
{   
    memset(task, 0, sizeof(struct task)); 
    task->page_directory = paging_new_4Gbyte(PAGING_IS_PRESENT | PAGING_ACCESS_FROM_ALL); 
    if (!task->page_directory) {
        return -TESTOS_STATUS_IO_ERROR; 
    }

    task->registers.ip = TESTOS_PROGRAM_VIRTUAL_ADDRESS;
    task->registers.cs = TESTOS_USER_CODE_SEGMENT; 
    task->registers.ss = TESTOS_USER_DATA_SEGMENT; 
    task->registers.esp = TESTOS_PROGRAM_VIRTUAL_ADDRESS_STACK_START; 
    task->process = process; 



    return 0; 
}

const int task_free(struct task* task) 
{
    paging_free_4Gbyte(task->page_directory); 
    task_list_remove(task); 
    kfree(task); 
    return 0; 
}

struct task* task_new(struct process* process)
{
    int res = 0; 
    struct task* task = kzalloc(sizeof(struct task)); 
    if (!task) {
        res = -TESTOS_STATUS_NOMEMORY; 
        goto out; 
    }

    res = task_init(task, process); 
    if (res != TESTOS_STATUS_OK) {
        goto out; 
    }

    if (!task_head) {
        task_head = task; 
        task_tail = task; 
        current_task = task; 
        goto out; 
    }

    task_tail->next = task;
    task->prev = task_tail; 
    task_tail = task; 

out:
    if (ISERR(res)) {
        task_free(task);
        return ERROR(res); 
    }

    return task; 
}

void task_save_state(struct task* task, struct interrupt_frame* frame)
{
    task->registers.ip = frame->ip; 
    task->registers.cs = frame->cs; 
    task->registers.flags = frame->flags; 
    task->registers.esp = frame->esp;
    task->registers.ss = frame->ss; 
    task->registers.eax = frame->eax;
    task->registers.ebp = frame->ebp;
    task->registers.ebx = frame->ebx;
    task->registers.ecx = frame->ecx;
    task->registers.edi = frame->edi;
    task->registers.edx = frame->edx;
    task->registers.esi = frame->esi;
}

void task_current_save_state(struct interrupt_frame* frame)
{
    if (!get_current_task()) {
        panic("task_current_save_state(): no current task to save");
    }

    struct task* task = get_current_task(); 
    task_save_state(task, frame); 
}

const int copy_string_from_task(struct task* task, void* virt, void* phys, const int max) 
{   
    if (max >= PAGING_PAGE_SIZE) {
        return -TESTOS_STATUS_INVALID_ARG; 
    }

    int res = 0; 
    char* tmp = kzalloc(max); 
    if (!tmp) {
        res = -TESTOS_STATUS_NOMEMORY; 
        return res; 
    }

    uint32_t* task_directory = task->page_directory->directory_entry; 
    uint32_t old_entry = paging_get(task_directory, tmp); 
    paging_map(task->page_directory, tmp, tmp, PAGING_IS_WRITEABLE | PAGING_IS_PRESENT | PAGING_ACCESS_FROM_ALL); 
    paging_switch(task->page_directory);
    strncpy(tmp, virt, max); 
    kernel_page(); 

    res = paging_set(task_directory, tmp, old_entry); 
    if (res < 0) {
        res = -TESTOS_STATUS_IO_ERROR; 
        goto out_free;
    }

    strncpy(phys, tmp, max); 

out_free:
    kfree(tmp); 

out:
    return res; 
}

void* task_get_stack_item(struct task* task, const int index) 
{
    void* res = 0; 
    uint32_t* stack_ptr = (uint32_t*)task->registers.esp; 

    // Switch to a given task page
    task_page_task(task); 

    res = (void*)stack_ptr[index]; 

    // Switch back to the kernel page
    kernel_page(); 

    return res; 
}